import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr

np.set_printoptions(precision=4, suppress=True)

def simulate_chain(P, start, n_steps):
    """
    Simulate Markov chain with transition matrix P starting from state index start
    for n_steps.
    Returns the list of visted states including the start state.
    :param P:
    :param start:
    :param n_steps:
    :return:
    """
    states = [start]
    cur = start
    for _ in range(n_steps):
        cur = npr.choice(a=len(P), p=P[cur])
        states.append(cur)

    return states

print("Section4: Periodic chain example (period cycle)")
P4 = np.array([
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
    [1.0, 0.0, 0.0]
])

print("Transition matrix P4:\n", P4)
print("This chain has period 3. We'll show that X_t returns to state 0 only at "
      " times multiple of 3.")

# simulate and record times of visits modulo 3
traj = simulate_chain(P4, 0, 3)
visit_times = [i for i,s in enumerate(traj) if s == 0]
print("Visit times to state 0 (including time 0):", visit_times)
